# -.- coding:utf-8 -.-
"""
Created on 2019-06-18 17:35
@author: cuizc
"""
import numpy as np
import align.detect_face
import tensorflow as tf
from facenet import facenet
from scipy.misc import imresize


def face_align():
    print('Creating networks and loading parameters')
    with tf.Graph().as_default():
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.5)
        sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, log_device_placement=True))
        pnet, rnet, onet = align.detect_face.create_mtcnn(sess, None)

        minsize = 20  # minimum size of face
        threshold = [0.6, 0.7, 0.7]  # three steps's threshold
        factor = 0.709  # scale factor

        def find_face(img):
            if img.ndim == 2:
                img = facenet.to_rgb(img)
            img = img[:, :, 0:3]

            bounding_boxes, _ = align.detect_face.detect_face(img, minsize, pnet, rnet, onet, threshold,
                                                              factor)
            nrof_faces = bounding_boxes.shape[0]
            if nrof_faces > 0:
                det = bounding_boxes[:, 0:4]
                det_arr = []
                img_size = np.asarray(img.shape)[0:2]
                for i in range(nrof_faces):
                    det_arr.append(np.squeeze(det[i]))
                positions = []
                faces = []
                for i, det in enumerate(det_arr):
                    det = np.squeeze(det)
                    bb = np.zeros(4, dtype=np.int32)
                    bb[0] = np.maximum(det[0] - 32 / 2, 0)
                    bb[1] = np.maximum(det[1] - 32 / 2, 0)
                    bb[2] = np.minimum(det[2] + 32 / 2, img_size[1])
                    bb[3] = np.minimum(det[3] + 32 / 2, img_size[0])
                    cropped = img[bb[1]:bb[3], bb[0]:bb[2], :]
                    scaled = imresize(cropped, (160, 160), interp='bilinear')
                    positions.append(bb)
                    faces.append(scaled)
                return positions, faces
            else:
                return None, None
        print("align model ready")
        return find_face
